//
//  MetalImageView.swift
//  ImageFilter
//
//  Created by Amos on 2025/10/10.
//

import SwiftUI
import Metal
import MetalKit

struct MetalImageView: UIViewRepresentable {
    // 输入图片
    var image: UIImage
    // 滤镜类型（可扩展）
    var filterType: FilterType
    // 亮度参数（仅亮度滤镜用）
    var brightness: Float
    
    // 滤镜类型枚举
    enum FilterType {
        case original
        case grayscale
        case brightness
    }
    
    func makeUIView(context: Context) -> MTKView {
        let mtkView = MTKView()
        mtkView.device = MetalContext.shared.device
        mtkView.clearColor = MTLClearColor(red: 0, green: 0, blue: 0, alpha: 1)
        mtkView.delegate = context.coordinator
        mtkView.drawableSize = image.size // 匹配图片尺寸
        return mtkView
    }
    
    func updateUIView(_ uiView: MTKView, context: Context) {
        // 更新滤镜参数时触发重绘
        context.coordinator.update(
            image: image,
            aFilterType: filterType,
            aBrightness: brightness
        )
        uiView.setNeedsDisplay()
    }
    
    func makeCoordinator() -> Coordinator {
        Coordinator(self)
    }
    
    // 协调器：处理 Metal 渲染逻辑
    class Coordinator: NSObject, MTKViewDelegate {
        private let parent: MetalImageView
        private let context = MetalContext.shared
        private var inputTexture: MTLTexture?
        private var pipelineState: MTLComputePipelineState?
        private var brightnessBuffer: MTLBuffer?
        
        // 滤镜类型（可扩展）
        private var filterType: FilterType?
        // 亮度参数（仅亮度滤镜用）
        private var brightness: Float?
        
        init(_ parent: MetalImageView) {
            self.parent = parent
            super.init()
            // 初始化纹理和管道
            update(image: parent.image, aFilterType: parent.filterType, aBrightness: parent.brightness)
        }
         
        // 更新输入数据和滤镜配置
        func update(image: UIImage, aFilterType: FilterType, aBrightness: Float) {
            
            filterType = aFilterType;
            brightness = aBrightness;
            
            // 转换 UIImage 为 Metal 纹理
            inputTexture = image.toMetalTexture(device: context.device)
            
            // 根据滤镜类型创建对应的计算管道
            let functionName: String
            switch filterType {
            case .original: functionName = "" // 原图无需处理
            case .grayscale: functionName = "grayscaleFilter"
            case .brightness: functionName = "brightnessFilter"
            case .none: functionName = ""
            }
            
            // 创建管道状态（原图模式跳过）
            if !functionName.isEmpty, let function = context.library.makeFunction(name: functionName) {
                pipelineState = try? context.device.makeComputePipelineState(function: function)
            } else {
                pipelineState = nil
            }
            var b:Float = aBrightness
            // 准备亮度参数缓冲区
            if filterType == .brightness {
                brightnessBuffer = context.device.makeBuffer(
                    bytes: &b,
                    length: MemoryLayout<Float>.stride
                )
            }
        }
        
        // MTKViewDelegate：绘制每一帧
        func draw(in view: MTKView) {
            guard let inputTexture = inputTexture,
                  let drawable = view.currentDrawable,
                  let commandBuffer = context.commandQueue.makeCommandBuffer() else { return }
            
            // 原图模式：直接将输入纹理渲染到视图
            if filterType == .original {
                renderOriginalTexture(inputTexture, to: drawable, commandBuffer: commandBuffer)
                return
            }
            
            // 滤镜模式：通过计算着色器处理
            guard let pipelineState = pipelineState else { return }
            
            // 创建输出纹理（与输入同尺寸）
            let outputTexture = createOutputTexture(from: inputTexture)
            
            // 编码计算命令
            guard let commandEncoder = commandBuffer.makeComputeCommandEncoder() else { return }
            commandEncoder.setComputePipelineState(pipelineState)
            commandEncoder.setTexture(inputTexture, index: 0)
            commandEncoder.setTexture(outputTexture, index: 1)
            
            // 设置亮度参数缓冲区（如果需要）
            if filterType == .brightness, let buffer = brightnessBuffer {
                commandEncoder.setBuffer(buffer, offset: 0, index: 0)
            }
            
            // 配置线程组
            let threadGroupSize = MTLSize(
                width: min(pipelineState.maxTotalThreadsPerThreadgroup, inputTexture.width),
                height: 1,
                depth: 1
            )
            let threadGroups = MTLSize(
                width: (inputTexture.width + threadGroupSize.width - 1) / threadGroupSize.width,
                height: inputTexture.height,
                depth: 1
            )
            commandEncoder.dispatchThreadgroups(threadGroups, threadsPerThreadgroup: threadGroupSize)
            commandEncoder.endEncoding()
            
            // 将处理后的纹理渲染到视图
            renderTexture(outputTexture, to: drawable, commandBuffer: commandBuffer)
        }
        
        // 渲染纹理到视图（使用简单的顶点着色器和片元着色器）
        private func renderTexture(_ texture: MTLTexture, to drawable: CAMetalDrawable, commandBuffer: MTLCommandBuffer) {
            // 简单的全屏渲染管道（固定顶点和纹理采样）
            
            // 1. 创建渲染通道描述符
            let renderPassDescriptor = MTLRenderPassDescriptor()
            
            // 2. 配置颜色附件（关键：设置渲染目标为 drawable 的纹理）
            let colorAttachment = renderPassDescriptor.colorAttachments[0]
            colorAttachment?.texture = drawable.texture  // 渲染目标：视图的可绘制纹理
            colorAttachment?.loadAction = .clear         // 绘制前清空
            colorAttachment?.clearColor = MTLClearColor(red: 0, green: 0, blue: 0, alpha: 1)  // 清空颜色
            colorAttachment?.storeAction = .store
            
            guard let encoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderPassDescriptor) else { return }
            
            encoder.setRenderPipelineState(context.pipelineState)
            encoder.setVertexBuffer(context.vertexBuffer, offset: 0, index: 0)
            encoder.setFragmentTexture(texture, index: 0)
            encoder.drawPrimitives(type: .triangleStrip, vertexStart: 0, vertexCount: 4) // 全屏矩形
            encoder.endEncoding()
            
            commandBuffer.present(drawable)
            commandBuffer.commit()
        }
        
        // 渲染原图（直接显示输入纹理）
        private func renderOriginalTexture(_ texture: MTLTexture, to drawable: CAMetalDrawable, commandBuffer: MTLCommandBuffer) {
            renderTexture(texture, to: drawable, commandBuffer: commandBuffer)
        }
        
        // 创建输出纹理
        private func createOutputTexture(from input: MTLTexture) -> MTLTexture {
            let descriptor = MTLTextureDescriptor()
            descriptor.pixelFormat = input.pixelFormat
            descriptor.width = input.width
            descriptor.height = input.height
            descriptor.usage = [.shaderRead, .shaderWrite, .renderTarget]
            return context.device.makeTexture(descriptor: descriptor)!
        }
        
        // MTKViewDelegate：视图尺寸变化时调用
        func mtkView(_ view: MTKView, drawableSizeWillChange size: CGSize) {}
    }
}

// 扩展：UIImage 转 MTLTexture
extension UIImage {
    func toMetalTexture(device: MTLDevice) -> MTLTexture? {
        guard let cgImage = cgImage else { return nil }
        let loader = MTKTextureLoader(device: device)
        return try? loader.newTexture(cgImage: cgImage, options: [.SRGB: false])
    }
}
